# coding:utf-8

"""
Q Learning 例子的 Maze(迷宫) 环境

黄色方形    机器人
红色方形    炸弹      [reward = -1]
绿色方形    宝藏      [reward = + 1]
其他方格    平地      [reward = 0]
"""

import sys
import time
import tkinter as tk
import numpy as np

# if sys.version_info.major == 2:
#     import Tkinte as tk
# else:
#     import tkinter as tk

# 迷宫的宽度
WIDTH = 15
# 迷宫的高度
HEIGHT = 15
# 每个方块的大小(像素值)
UNIT = 40


class Maze(tk.Tk, object):
    def __init__(self):
        super(Maze, self).__init__()
        # 上下左右四个action(动作)
        self.action_space = ["u", "d", "l", "r"]
        # action的数目
        self.n_actions = len(self.action_space)
        # 游戏窗口标题
        self.title("Q Learning")
        # Tkinter的几何形状
        self.geometry("{0}x{1}".format(WIDTH * UNIT, HEIGHT * UNIT))
        self.build_maze()

    # 创建迷宫
    def build_maze(self):
        # 创建画布 Canvas
        self.canvas = tk.Canvas(self, bg="white", width=WIDTH * UNIT, height=HEIGHT * UNIT)
        # 绘制横纵放个线
        for c in range(0, WIDTH * UNIT, UNIT):
            x0, y0, x1, y1 = c, 0, c, HEIGHT * UNIT
            self.canvas.create_line(x0, y0, x1, y1)
        for r in range(0, HEIGHT * UNIT, UNIT):
            x0, y0, x1, y1 = 0, r, WIDTH * UNIT, r
            self.canvas.create_line(x0, y0, x1, y1)

        # 零点(左上角)
        # 方格大小是 40 × 40  ， 20,20 就是中心的位置
        origin = np.array([20, 20])

        # 创建我们的探索者 机器人(robot)
        robot_center = origin + np.array([0, UNIT * 2])
        self.robot = self.canvas.create_oval(
            robot_center[0] - 15, robot_center[1] - 15,
            robot_center[0] + 15, robot_center[1] + 15,
            fill="yellow"
        )

        # 炸弹1
        bomb1_center = origin + UNIT
        self.bomb1 = self.canvas.create_rectangle(
            bomb1_center[0] - 15, bomb1_center[1] - 15,
            bomb1_center[0] + 15, bomb1_center[1] + 15,
            fill="red"
        )

        # 炸弹2
        bomb2_center = origin + np.array([UNIT * 3, UNIT])
        self.bomb2 = self.canvas.create_rectangle(
            bomb2_center[0] - 15, bomb2_center[1] - 15,
            bomb2_center[0] + 15, bomb2_center[1] + 15,
            fill="red"
        )

        # 宝藏
        treasure_center = origin + np.array([UNIT * 3, 0])
        self.treasure = self.canvas.create_rectangle(
            treasure_center[0] - 15, treasure_center[1] - 15,
            treasure_center[0] + 15, treasure_center[1] + 15,
            fill="green"
        )

        # 设置好上面配置的场景
        self.canvas.pack()

    # 重置(游戏重新开始，将机器人放到左下角)
    def reset(self):
        self.update()
        time.sleep(0.5)
        # 删去机器人
        self.canvas.delete(self.robot)
        origin = np.array([20, 20])
        robot_center = origin + np.array([0, UNIT * 2])

        # 重新创建机器人
        self.robot = self.canvas.create_oval(
            robot_center[0] - 15, robot_center[1] - 15,
            robot_center[0] + 15, robot_center[1] + 15,
            fill="yellow"
        )

        # 返回 观测(observation)
        return self.canvas.coords(self.robot)

    # 走一步 (机器人实施 action)
    def step(self, action):
        s = self.canvas.coords(self.robot)
        base_action = np.array([0, 0])
        # 上
        if action == 0:
            if s[1] > UNIT:
                base_action[1] -= UNIT
        # 下
        elif action == 1:
            if s[1] < (HEIGHT - 1) * UNIT:
                base_action[1] += UNIT
        # 右
        elif action == 2:
            if s[0] < (WIDTH - 1) * UNIT:
                base_action[0] += UNIT
        # 左
        elif action == 3:
            if s[0] > UNIT:
                base_action[0] -= UNIT

        # 移动机器人
        self.canvas.move(self.robot, base_action[0], base_action[1])

        # 下一个 state
        s_ = self.canvas.coords(self.robot)

        # 奖励机制
        if s_ == self.canvas.coords(self.treasure):
            # 找到宝藏 奖励为 1
            reward = 1
            done = True
            # 终止
            s_ = "terminal"
            print("找到宝藏， 好棒！")
        elif s_ == self.canvas.coords(self.bomb1):
            # 找到炸弹1 奖励为 -1
            reward = -1
            done = True
            # 终止
            s_ = "terminal"
            print("找到炸弹1， GG！")

        elif s_ == self.canvas.coords(self.bomb2):
            # 找到炸弹2 奖励为1
            reward = -1
            done = True
            # 终止
            s_ = "terminal"
            print("找到炸弹2， GG！")
        else:
            # 其他格子 继续移动 没有奖励
            reward = 0
            done = False

        return s_, reward, done

    # 调用 Tkinter 的 update 方法
    def render(self):
        time.sleep(0.1)
        self.update()
